#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#define min(X, Y)  ((X) < (Y) ? (X) : (Y))
#define max(X, Y)  ((X) > (Y) ? (X) : (Y))
#define INF 1000

int myround(float r){
  if (r-floorf(r)<0.5) return (int)floorf(r);
  return (int)ceilf(r);
}

float lnfact(int n){
  static float table[1000];
  return (table[n]>0)?table[n]:(table[n]=lgamma(n+1)); 
}

float lnbinomial(int n, int k){
  static float table[100*100];
  if (n<0 || k<0)  {return -INF;}
  if (k>n) {return -INF;}
  //return lnfact(n)-lnfact(k)-lnfact(n-k);
  return (table[n*100+k]>0)?table[n*100+k]:(table[n*100+k]=lnfact(n)-lnfact(k)-lnfact(n-k)); 
}
float lnmultinomial(int n, int k1, int k2){
  if (n-k1-k2<0  ) {return -INF;}
  if (n <0 || k1<0 || k2<0) {return -INF;}
  return lnfact(n)-lnfact(k1)-lnfact(k2)-lnfact(n-k1-k2);
}
float a(int n){
  if (0==n) {fprintf(stderr,"a:invalid n\n");exit(1);}
  int i;
  float s=0.0;
  for (i=1;i<n;i++){
    s+=1.0/i;
  }
  return s;
}
float beta(int n,int i){
   int d1=n-i;
   if (d1==0) return 0;
   return 2.0*n*(a(n+1)-a(i))/(d1*(d1+1))-2.0/d1; 
}

float ps(int k, int l, int n, float theta){
  int mkl=min(k,l);
  float common_term=theta*theta*(beta(n,mkl)-beta(n,mkl+1))/2.0;
  float first_term=0.0;
  if (k==l) {
    return theta*theta*beta(n,k);
  }
  return common_term;
}
float pe(int k, int l, int n, float theta){
  float res=0.0;
  if (k+l>n) return res;
  if (k+l==n){
    res=theta*theta*(((a(n)-a(k))/(float)(n-k))+((a(n)-a(l))/(float)(n-l))-((beta(n,k)+beta(n,l))/2.0));
  } else{
    res=theta*theta*(1.0/(k*l) - (beta(n,k)-beta(n,k+1)+beta(n,l)-beta(n,l+1))/2.0 );
    //printf("res:%f, %f\n",res,(beta(n,k)-beta(n,k+1)+beta(n,l)-beta(n,l+1)));
  }
  return res;
}
float sigma(int i, int j,int n){
  int tmp;
  if (i<j) {tmp=i;i=j;j=tmp;}
  if (i==j){ 
    if ( 2*i < n) return beta(n,i+1);
    if ( 2*i == n ) return 2.0*(a(n)-a(i))/(n-i) - 1.0/(i*i);
    return beta(n,i)-1.0/(i*i);
  }
  if (i+j<n) return (beta(n,i+1)-beta(n,i))/2.0;
  if (i+j==n) return  ((a(n)-a(i))/(n-i))+((a(n)-a(j))/(n-j))-(beta(n,i)+beta(n,j+1))/2.0-1.0/(i*j);
  return (beta(n,j)-beta(n,j+1))/2.0-1.0/(i*j);
}
float fu_sum(int k, int l, int n, float theta){
  return theta*theta*(1.0/(k*l)+sigma(k,l,n));
}
int min3(int a, int b, int c){
  return min(min(a,b),c);
}
int max3(int a, int b, int c){
  return max(max(a,b),c);
}
float cs(int i, int j, int k, int l, int nx, int ny, int nxy,float theta){
  if ((l-j<0)||(nx-nxy-l+j<0)||(k-i<0)||(ny-nxy-k+i<0)||(nx+ny-nxy-k<0)) return 0;
  float first_factor = exp(lnbinomial(nx-nxy,l-j)+lnbinomial(ny - nxy,k - i)- lnmultinomial(nx+ny-nxy,l,k-l));
  int lb = max(i-nxy,l-j);
  int ub = min3(i,nx-nxy,k-j); 
  int kx;
  float sum=0.0;
  for(kx=lb;kx<=ub;kx++){
    float b1=lnbinomial(nxy,i-kx);
    float b2=lnbinomial(nx-nxy+j-l,kx+j-l);
    float b3=lnbinomial(k-kx,j);
    sum+=exp(b1+b2+b3); 
  }
  return first_factor*sum;
}
float ce(int i, int j, int k, int l, int nx, int ny, int nxy,float theta){
  if ((l-j<0)||(nx-nxy-l+j<0)||(k-i<0)||(ny-nxy-k+i<0)||(nx+ny-nxy-k-l<0)) return 0.0;
  if (k+l>nx+ny-nxy) return 0.0;
  float first_factor = exp(lnbinomial(nx-nxy,l-j)+lnbinomial(ny - nxy,k - i)- lnmultinomial(nx+ny-nxy,k,l)); 
  //fprintf(stdout,"i,j,k,l:%d,%d,%d,%d first factor:%g\n",i,j,k,l,first_factor);
  int lb = max3(0,i-nxy,k-ny+j);
  int ub = min(i,nx-nxy+j-l);
  int kx;
  float sum=0.0;
  //fprintf(stdout,"lb=%d,ub=%d\n",lb,ub);
  for(kx=lb;kx<=ub;kx++){
    float b1=lnbinomial(nxy,i-kx);
    float b2=lnbinomial(nx-nxy+j-l,kx);
    float b3=lnbinomial(ny-k+kx,j);
    //printf("b1=%f,b2=%f,b3=%f\n",exp(b1),exp(b2),exp(b3));
    sum+=exp(b1+b2+b3); 
  }
  return first_factor*sum;
}
float p(int i, int j, int nx, int ny, int nxy, float theta){
  int lb = 1;
  int ub = nx +ny -nxy -1 ;
  float sum = 0.0;
  float term;
  int k,l;
  float c_s,p_s,c_e,p_e;
  for (l=lb;l<=ub;l++){
    for(k=lb;k<=ub;k++){
      if (k>=l){
        c_s=cs(i,j,k,l,nx,ny,nxy,theta);
        p_s=ps(k,l,nx+ny-nxy,theta);
        c_e=ce(i,j,k,l,nx,ny,nxy,theta);
        p_e=pe(k,l,nx+ny-nxy,theta);
        term=c_s*p_s+c_e*p_e;
        }
      else{
        term=cs(j,i,l,k,ny,nx,nxy,theta)*ps(k,l,nx+ny-nxy,theta)+ce(j,i,l,k,ny,nx,nxy,theta)*pe(k,l,nx+ny-nxy,theta);
      }  
    //if (isnan(term)) {printf("%d %d %d %d\n",i,j,k,l);} 
    sum+=term;
    }
  }
  return sum;  
}
float cov(int i, int j, int nx, int ny, int nxy, float theta){
  return p(i,j,nx,ny,nxy,theta)- theta*theta/(i*j);
}
void generate3vectors(int l, float pm, int n, int* nx, int* ny, int* nxy ){
  int i,j;
  float p1,p2;
  for(i=0;i<l;i++){
    for (j=0;j<n;j++){
      p1=rand()/(float)RAND_MAX;p2=rand()/(float)RAND_MAX;
      //printf("p1=%f,p2=%f\n",p1,p2);
      if (p1>pm) {nx[i]++;}
      if (p2>pm) {ny[i]++;}
      if (p1>pm && p2>pm) {nxy[i]++;}
    }
    if (nx[i]<2) nx[i]=2;
    if (ny[i]<2) ny[i]=2;
  }
}
void generate3vectorsplusdata(int l, float pm, int n, int* nx, int* ny, int* nxy, int* allele, float* t, float* d, float* h, float theta, float ratio, float freq){
  int i,j,ax;
  float p1,p2,p3;
  float asum=0.0;  
  for(i=0;i<l;i++){
    //ax=0;
    for (j=0;j<n;j++){
      p1=rand()/(float)RAND_MAX;p2=rand()/(float)RAND_MAX;
      //printf("p1=%f,p2=%f\n",p1,p2);
      if (p1>pm) {nx[i]++;}// p3=rand()/(float)RAND_MAX; if (p3<(float)(allele[i])/(float)n) ax++;}
      if (p2>pm) {ny[i]++;}
      if (p1>pm && p2>pm) {nxy[i]++;}
    }
    if (nx[i]<2) {nx[i]=2;}// ax=0; p3=rand()/(float)RAND_MAX; if (p3<(float)(allele[i])/(float)n) ax++; p3=rand()/(float)RAND_MAX; if (p3<(float)(allele[i])/(float)n) ax++;}
    if (ny[i]<2) ny[i]=2;
  }
  for(i=0;i<l;i++){
    asum+=a(nx[i]);
  }
  for(i=0;i<l;i++){  
    //if(ax>0 && ax<nx[i]) t[i]=1; else t[i]=0;
    p3=rand()/(float)RAND_MAX;
    //for(allele[i]=1;(p3>0 && allele[i]<nx[i]);allele[i]++){ p3+=-theta/(float)(allele[i]); }
    if(p3>0) { if(p3<theta*ratio) {p1=rand()/(float)RAND_MAX;
    allele[i]=myround(0.5*(1.0+p1)*(float)nx[i]);} else allele[i]=0; } else {allele[i]-=1;}
    if (allele[i]==0 || allele[i]==n) t[i]=0; else {
      t[i]=1;
    }
    d[i]=(float)(2*allele[i]*(nx[i]-allele[i]))/(float)(l*nx[i]*(nx[i]-1));
    h[i]=-(float)(allele[i])/(float)(l*(nx[i]-1));
    //printf("allele %d %d %d %g %g\n",allele[i],n,nx[i],t[i],d[i]);    //debug    
  }
  for(i=0;i<l;i++){
    t[i]=t[i]/asum;
    h[i]=d[i]-h[i];
    d[i]=d[i]-t[i];
  }
}
float watterson_variance(float theta, int l, int* nx, int* ny,int* nxy, float* var_d, float* var_h){
  float asum=0.0;
  int x,y;
  for(x=0;x<l;x++){
    asum+=a(nx[x]);
  }
  float covsum=0.0;
  float checksum=0.0;
  float dsum=0.0;
  float covdsum=0.0;
  float checksum_h=0.0;
  float hsum=0.0;
  float covhsum=0.0;
  int i,j;
  float wxd,wyd,wxh,wyh;
  float tmpcov;
  for(x=0;x<l;x++){
      printf("x=%d\r",x);
      checksum=0.0;
      checksum_h=0.0;
      for(i=1;i<=nx[x]-1;i++){
	//wx=(float)(2*i*(nx[x]-i))/(float)(l*nx[x]*(nx[x]-1))-1.0/asum;
        wxh=(float)(2*i*(nx[x]-i))/(float)(l*nx[x]*(nx[x]-1))-(float)i/(float)(l*(nx[x]-1));
        wxd=(float)(2*i*(nx[x]-i))/(float)(l*nx[x]*(nx[x]-1))-1.0/asum;
	//printf("w[%d,%d]=%f\n",i,nx[x],wx);
	checksum+=wxd/(float)i;//debug
	dsum+=wxd*wxd/(float)i;
        checksum_h+=wxh/(float)i;//debug
	hsum+=wxh*wxh/(float)i;
        for(j=1;j<=ny[x]-1;j++){
          wyh=(float)(2*j*(ny[x]-j))/(float)(l*ny[x]*(ny[x]-1))-(float)j/(float)(l*(ny[x]-1));
          wyd=(float)(2*j*(ny[x]-j))/(float)(l*ny[x]*(ny[x]-1))-1.0/asum;
          tmpcov=cov(i,j,nx[x],ny[x],nxy[x],theta);
	  covsum+=tmpcov;
	  covdsum+=wxd*wyd*tmpcov;
          covhsum+=wxh*wyh*tmpcov;
        }
      }
    }
  *var_d=dsum*theta+covdsum*(float)l;
  *var_h=hsum*theta+covhsum*(float)l;
  return (theta/asum)+(1.0/(asum*asum))*covsum*l;     
}
int main(int argc, char* argv[]){
  int k,l,n,i,j,nx,ny,nxy;
  float theta;
  k=atoi(argv[1]);
  l=atoi(argv[2]);
  n=atoi(argv[3]);
  theta=atof(argv[4]); 
  i=atoi(argv[5]);
  j=atoi(argv[6]);
  nx=atoi(argv[7]);
  ny=atoi(argv[8]);
  nxy=atoi(argv[9]); 
  printf("factorial of %d is:%f\n",n,exp(lnfact(n)));
  printf("%d over %d = %f\n",n,k,exp(lnbinomial(n,k)));
  printf("a (n=%d)=%f\n",n,a(n));
  printf("beta(n=%d,k=%d)=%f\n",n,k,beta(n,k));
  printf("sigma (k=%d,l=%d,n=%d)=%g\n",k,l,n,sigma(k,l,n));
  printf("pE+pS computed with Fu[95]=%g\n",fu_sum(k,l,n,theta));
  printf("pE=%g\n",pe(k,l,n,theta));
  printf("pS=%g\n",ps(k,l,n,theta));
  printf("pE+pS=%g\n",pe(k,l,n,theta)+ps(k,l,n,theta));
  printf("i=%d,j=%d,nx=%d,ny=%d,nxy=%d\n",i,j,nx,ny,nxy);
  printf("C^E=%g\n",ce(i,j,k,l,nx,ny,nxy,theta));
  printf("C^S=%g\n",cs(i,j,k,l,nx,ny,nxy,theta));
  printf("p=%g\n",p(i,j,nx,ny,nxy,theta));
  //printf("cov=%f\n",cov(i,j,nx,ny,nxy,theta));
  int len=atoi(argv[10]);
  float pm=atof(argv[11]); 
  float var_w,var_d,var_h;
  int* nxv = malloc(len*sizeof(int));
  int* nyv = malloc(len*sizeof(int));
  int* nxyv = malloc(len*sizeof(int));
  for(i=0;i<len;i++){
    nxv[i]=0;nyv[i]=0;nxyv[i]=0;
  }
  generate3vectors(len,pm,n,nxv,nyv,nxyv);
  FILE * input;
  input=fopen("sequences.txt","r");
  printf("len=%d\n",len);
  for(i=0;i<len;i++){
    fprintf(stdout,"pos %d nx=%d,ny=%d,nxy=%d\n",i,nxv[i],nyv[i],nxyv[i]);
  }
  //printf("variance of watterson:%f\n",watterson_variance(theta,len,nxv,nyv,nxyv,&var_d));
  printf("binomial of 10,6:%f\n",exp(lnbinomial(10,6)));
  printf("multinomial of 10 2 3:%f\n",exp(lnmultinomial(10,2,3)));
  float s=0; 
  for(i=1;i<n;i++){
    for(j=1;j<n;j++){
      //printf("ecco qua (porca puzzola):%d %d %g\n",i,j,theta*theta*sigma(i,j,n));//debug
      s+=theta*theta*sigma(i,j,n);
    }
  }
  printf("s=%g\n",s);
  
  int nf=atof(argv[12]);
  int sim;
  int nsims=atof(argv[13]);
  int *allele=malloc(len*sizeof(int));
  float sum_t, sum_d, sum_h;
  float* value_t = malloc(len*sizeof(float));
  float* value_d = malloc(len*sizeof(float));
  float* value_h = malloc(len*sizeof(float));
  //generatefulldata(len,n,theta,0.5,4.0,input,allele);
  printf("Data of Watterson and Tajima variance for fixed sample size %d:\n",n);
  for(j=0;j<=nf-1;j++){
    for(i=0;i<len;i++){
      nxv[i]=0;nyv[i]=0;nxyv[i]=0;
    }
    sum_t=0;
    sum_d=0;
    sum_h=0;
    for(sim=1;sim<=nsims;sim++){
      //printf("sim %d\n",sim);
      for(i=0;i<len;i++){
        nxv[i]=0;nyv[i]=0;nxyv[i]=0;
      }
      generate3vectorsplusdata(len,(float)j/(float)nf,n,nxv,nyv,nxyv,allele,value_t,value_d,value_h,theta,4.0,0.5);
      //printf("len=%d\n",len);
      //for(i=0;i<len;i++){
	//fprintf(stdout,"pos %d nx=%d,ny=%d,nxy=%d\n",i,nxv[i],nyv[i],nxyv[i]);
      //}    
      for(i=0;i<len;i++){
	sum_t+=value_t[i];
	sum_d+=value_d[i];
	sum_h+=value_h[i];
      }
    }
    sum_t=sum_t/(float)nsims;
    sum_d=sum_d/(float)nsims;
    sum_h=sum_h/(float)nsims;
    var_w=watterson_variance(theta,len,nxv,nyv,nxyv,&var_d,&var_h);
    printf("%f %d %g %g %g %g %g %g\n",(float)j/(float)nf,n,sum_t,var_w,sum_d,var_d,sum_h,var_h);
  }
  printf("Data of Watterson and Tajima variance for equivalent sample size=coverage : \n",n);
  for(j=0;j<=nf-1;j++){
    for(i=0;i<len;i++){
      nxv[i]=0;nyv[i]=0;nxyv[i]=0;
    }
    sum_t=0;
    sum_d=0;
    sum_h=0;
    for(sim=1;sim<=nsims;sim++){
      //printf("sim %d\n",sim);
      for(i=0;i<len;i++){
        nxv[i]=0;nyv[i]=0;nxyv[i]=0;
      }
      generate3vectorsplusdata(len,0.0,myround((float)(n*(nf-j))/(float)nf),nxv,nyv,nxyv,allele,value_t,value_d,value_h,theta,4.0,0.5);
      //printf("len=%d\n",len);
      //for(i=0;i<len;i++){
	//fprintf(stdout,"pos %d nx=%d,ny=%d,nxy=%d\n",i,nxv[i],nyv[i],nxyv[i]);
      //}    
      for(i=0;i<len;i++){
	sum_t+=value_t[i];
	sum_d+=value_d[i];
	sum_h+=value_h[i];
      }
    }
    sum_t=sum_t/(float)nsims;
    sum_d=sum_d/(float)nsims;
    sum_h=sum_h/(float)nsims;
    var_w=watterson_variance(theta,len,nxv,nyv,nxyv,&var_d,&var_h);
    printf("%f %d %g %g %g %g %g %g\n",(float)j/(float)nf,n,sum_t,var_w,sum_d,var_d,sum_h,var_h);
  }
  printf("Data of Watterson and Tajima variance for fixed coverage %d:\n",n);
  for(j=0;j<=nf-1;j++){
    for(i=0;i<len;i++){
      nxv[i]=0;nyv[i]=0;nxyv[i]=0;
    }
    sum_t=0;
    sum_d=0;
    sum_h=0;
    for(sim=1;sim<=nsims;sim++){
      //printf("sim %d\n",sim);
      for(i=0;i<len;i++){
        nxv[i]=0;nyv[i]=0;nxyv[i]=0;
      }
      generate3vectorsplusdata(len,(float)j/(float)nf,myround((float)(n*nf)/(float)(nf-j)),nxv,nyv,nxyv,allele,value_t,value_d,value_h,theta,4.0,0.5);
      //printf("len=%d\n",len);
      //for(i=0;i<len;i++){
	//fprintf(stdout,"pos %d nx=%d,ny=%d,nxy=%d\n",i,nxv[i],nyv[i],nxyv[i]);
      //}    
      for(i=0;i<len;i++){
	sum_t+=value_t[i];
	sum_d+=value_d[i];
	sum_h+=value_h[i];
      }
    }
    sum_t=sum_t/(float)nsims;
    sum_d=sum_d/(float)nsims;
    sum_h=sum_h/(float)nsims;
    var_w=watterson_variance(theta,len,nxv,nyv,nxyv,&var_d,&var_h);
    printf("%f %d %g %g %g %g %g %g\n",(float)j/(float)nf,myround((float)(n*nf)/(float)(nf-j)),sum_t,var_w,sum_d,var_d,sum_h,var_h);
  }
  
  //fclose(input);
  
}
